package cn.itcast.flink.trans;

import org.apache.commons.lang3.StringUtils;
import org.apache.flink.api.common.functions.Partitioner;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

/**
 * Author itcast
 * Date 2022/1/13 10:43
 * Desc 彩票案例
 * 开发步骤:
 * 1. 生成环境
 * 2. 设置并行度
 * 3. 生成彩票类
 * 4. 根据彩票id进行分区
 * 5. 统计
 * 6. 打印输出
 * 7. 执行流环境
 */
public class PartitionCustomDemo {
    public static void main(String[] args) throws Exception {
        //1. 生成环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        //2. 设置并行度
        env.setParallelism(2);
        //3. 生成彩票类
        DataStreamSource<Caipiao> source = env.fromElements(
                new Caipiao("12", "2021-12-20 10:10:10", 20.00),
                new Caipiao("10", "2021-12-20 10:10:10", 20.00),
                new Caipiao("21", "2021-12-20 10:10:10", 20.00),
                new Caipiao("22", "2021-12-20 10:10:10", 20.00),
                new Caipiao("23", "2021-12-20 10:10:10", 20.00),
                new Caipiao("1", "2021-12-20 10:10:10", 20.00)

        );
        //4. 根据彩票id进行分区
        DataStream<Caipiao> result = source.partitionCustom(
                new Partitioner<String>() {
                    @Override
                    public int partition(String key, int numPartitions) {
                        if (StringUtils.isNotEmpty(key)) {
                            int caipiaoId = Integer.parseInt(key);
                            if (caipiaoId > 15) {
                                return 0;
                            } else {
                                return 1;
                            }
                        } else {
                            return 0;
                        }
                    }
                },
                new KeySelector<Caipiao, String>() {
                    @Override
                    public String getKey(Caipiao value) throws Exception {
                        return value.getCid();
                    }
                }
        );
        //5. 统计
        //6. 打印输出
        result.print();
        //7. 执行流环境
        env.execute();
    }


    public static class Caipiao {
        private String cid;
        private String startTime;
        private double money;

        public Caipiao(String cid, String startTime, double money) {
            this.cid = cid;
            this.startTime = startTime;
            this.money = money;
        }

        public String getCid() {
            return cid;
        }

        public void setCid(String cid) {
            this.cid = cid;
        }

        public String getStartTime() {
            return startTime;
        }

        public void setStartTime(String startTime) {
            this.startTime = startTime;
        }

        public double getMoney() {
            return money;
        }

        public void setMoney(double money) {
            this.money = money;
        }

        @Override
        public String toString() {
            return "Caipiao{" +
                    "cid='" + cid + '\'' +
                    ", startTime='" + startTime + '\'' +
                    ", money=" + money +
                    '}';
        }
    }
}
